#include "iostream"
#include "opencv2/opencv.hpp"
using namespace std;

size_t clip(size_t value, size_t min_v, size_t max_v) {
    if (value < min_v)
        value = min_v;
    if (value > max_v)
        value = max_v;
    return value;
}

void cal_gray_saliency(cv::Mat &image, cv::Mat &out, size_t size, int k) {
    size_t radius = (size - 1) / 2;
    int height = image.size().height, width = image.size().width;
    double H[k][k], P[k][k];
    size_t H_sum = 0, P_non_zero_sum = 0;

    for (int i = 0; i < k; ++ i) {
        for (int j = 0; j < k; ++ j) {
            H[i][j] = 0;
        }
    }

    int center_p, neigh_p, locx, locy;
    for (int i = 0; i < height; ++ i) {
        for (int j = 0; j < width; ++ j) {
            center_p = image.at<uint8_t>(i, j);
            for (int wi = 0; wi < size; ++ wi) {
                for (int wj = 0; wj < size; ++ wj) {
                    locx = i + wi - radius;
                    locy = j + wj - radius;
                    if (
                        locx >= 0 && locx <= height - 1 &&
                        locy >= 0 && locy <= width - 1
                    ) {
                        // locx = clip(i + wi - radius, 0, height - 1);
                        // locy = clip(j + wj - radius, 0, width - 1);
                        neigh_p = image.at<uint8_t>(locx, locy);
                        if (neigh_p != center_p) {
                            H[center_p][neigh_p] ++;
                            H[neigh_p][center_p] ++;
                            H_sum ++;
                        }
                    }
      
                }
            }
        }
    }

    double double_H_sum = (double)H_sum, p_value;
    for (int i = 0; i < k; ++ i) {
        for (int j = 0; j < k; ++ j) {
            P[i][j] = H[i][j] / double_H_sum;
            if (P[i][j] != 0) {
                P_non_zero_sum ++;
            }
        }
    }
    double double_P_non_zero_sum = (double)P_non_zero_sum;
    double U = 1 / double_P_non_zero_sum;

    cv::Mat S = cv::Mat::zeros(height, width, CV_64F);

    for (int i = 0; i < height; ++ i) {
        for (int j = 0; j < width; ++ j) {
            center_p = image.at<uint8_t>(i, j);
            for (size_t wi = 0; wi < size; ++ wi) {
                for (size_t wj = 0; wj < size; ++ wj) {
                    locx = i + wi - radius;
                    locy = j + wj - radius;
                    if (
                        locx >= 0 && locx <= height - 1 &&
                        locy >= 0 && locy <= width - 1
                    ) {
                        // locx = clip(i + wi - radius, 0, height - 1);
                        // locy = clip(j + wj - radius, 0, width - 1);
                        neigh_p = image.at<uint8_t>(locx, locy);
                        p_value = P[center_p][neigh_p];
                        if (p_value > U) {
                            p_value = 0;
                        } else if(p_value > 0) {
                            p_value = U - p_value;
                        }
                        S.at<double>(i, j) += p_value;
                    }
                }
            }
        }
    }

    cv::Mat Test = cv::Mat::zeros(k, k, CV_64F);
    for (int i = 0; i < k; ++ i) {
        for (int j = 0; j < k; ++ j) {
            Test.at<double>(i, j) = P[i][j];
        }
    }
    
    cv::normalize(S, out, 1.0, 0.0, cv::NORM_MINMAX);

    
    // double max_value = 0;
    // for (int i = 0; i < height; ++ i) {
    //     for (int j = 0; j < width; ++ j) {
    //         max_value = max(max_value, S.at<double>(i, j));
    //     }
    // }
    // for (int i = 0; i < height; ++ i) {
    //     for (int j = 0; j < width; ++ j) {
    //         out.at<double>(i, j) = S.at<double>(i, j) / max_value;
    //     }
    // }

}

void saliency(cv::Mat &in, cv::Mat &out, int size, int guassion_size, double sigmaX, int use_grad) {
    int channel = in.channels(), height = in.size().height, width = in.size().width;
    cv::Mat channels[3];
    cv::split(in, channels);
	cv::Mat sobelx, sobely, angle_map, gray_saliency = cv::Mat::zeros(height, width, CV_64F), S = cv::Mat::zeros(height, width, CV_64F);

    for (int c = 0; c < channel; ++ c) {

		cal_gray_saliency(channels[c], gray_saliency, size, 256);
		cv::add(S, gray_saliency, S);

        if (use_grad) {
            cv::Sobel(channels[c], sobelx, CV_32F, 1, 0, 3);
            cv::Sobel(channels[c], sobely, CV_32F, 0, 1, 3);
            cv::phase(sobelx, sobely, angle_map, true);
            angle_map.convertTo(angle_map, CV_8U, 0.5);
            cal_gray_saliency(angle_map, gray_saliency, size, 181);
            cv::add(S, gray_saliency, S);
        }

        cv::GaussianBlur(S, S, {guassion_size, guassion_size}, sigmaX);
    }	

	cv::normalize(S, S, 1.0, 0.0, cv::NORM_MINMAX);
    S.convertTo(out, CV_8U, 255);
}

int main(int argc, char const *argv[])
{
	if (argc != 6) {
        cout << "useage: "<<  argv[0] << " <input_image> <output_image> size guassion_size sigmaX" << endl;
		exit(-1);
	}
    cv::Mat image = cv::imread(argv[1]);    
    cv::Mat output;
    cv::cvtColor(image, image, cv::COLOR_BGR2Lab);
    saliency(image, output, atoi(argv[3]), atoi(argv[4]), atof(argv[5]), 1);
    cv::imwrite(argv[2], output);

    return 0;
}